use crate::{click_set::ClickSet, coordinate::Coord, state::State, strategy::Strategy};
use rand::prelude::*;
use std::cmp::Reverse;

pub struct World<'a> {
    state: &'a State,
    /// 非空，按分数从高到低排序
    strategies: Vec<Strategy>,
}

impl<'a> World<'a> {
    pub fn new(state: &'a State, size: usize, max_size: usize, step: usize) -> Self {
        let strategy = {
            let mut click_set = ClickSet::new((**state).clone());
            click_set.clicks_mut()[0] += step;
            Strategy::new(click_set, &mut state.clone())
        };
        let mut strategies = Vec::with_capacity(max_size);
        let stream = itertools::iterate(strategy, |s| modified(s, &mut state.clone()));
        strategies.extend(stream.take(size));
        strategies.sort_unstable_by_key(|s| Reverse(s.score));
        Self { state, strategies }
    }

    fn reproduce(&mut self) {
        let mut rng = thread_rng();
        while self.strategies.len() < self.strategies.capacity() {
            self.strategies.push(modified(
                self.strategies.choose(&mut rng).expect("没有策略"),
                &mut self.state.clone(),
            ));
        }
    }

    fn fittest(&mut self, keep: usize) {
        use rand::seq::index::sample_weighted;
        self.strategies.sort_unstable_by_key(|s| Reverse(s.score));
        let dead = sample_weighted(
            &mut thread_rng(),
            self.strategies.len(),
            |i| i as f64,
            self.strategies.len() - keep,
        )
        .expect("选择错误");
        for index in dead {
            self.strategies[index].score = 0; // 假定没有零分
        }
        self.strategies.retain(|s| s.score != 0);
    }

    pub fn run(&mut self) {
        let keep = self.strategies.len();
        self.reproduce();
        self.fittest(keep);
    }

    pub fn best(&self) -> &Strategy {
        self.strategies.first().expect("没有策略")
    }
}

fn modified(s: &Strategy, state: &mut State) -> Strategy {
    let mut rng = thread_rng();
    let mut click_set = (*s).clone();
    let mut src = [Coord::default(); 3];
    let src = {
        let len = click_set
            .iter_cells()
            .filter(|&coord| click_set[coord] > 0)
            .choose_multiple_fill(&mut rng, &mut src);
        &src[0..len]
    };
    let mut dest = [Coord::default(); 3];
    let dest = {
        let len = click_set
            .iter_cells()
            .choose_multiple_fill(&mut rng, &mut dest);
        &dest[0..len]
    };
    for (&src, &dest) in std::iter::zip(src, dest) {
        click_set[src] -= 1;
        click_set[dest] += 1;
    }
    Strategy::new(click_set, state)
}
